package common.util;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.ReflectUtil;

import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.Optional;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.function.Predicate;
import java.util.function.Supplier;
import java.util.stream.Collector;

/**
 * 树工具类
 *
 * @author 米泽鹏
 * @since 2023-11-24 下午 11:18
 */
public class TreeUtil extends cn.hutool.core.lang.tree.TreeUtil {

	/**
	 * 树形结构迭代器（向下，深度优先）
	 *
	 * @param roots             根节点集合
	 * @param handler           处理器
	 * @param childrenFieldName 子节点字段名
	 * @param <T>               节点类型
	 */
	@SuppressWarnings("unchecked")
	public static <T> void treeIteratorDown(Collection<T> roots, Consumer<T> handler, String childrenFieldName) {
		if (CollUtil.isEmpty(roots)) {
			return;
		}
		for (T root : roots) {
			handler.accept(root);
			treeIteratorDown((Collection<T>) ReflectUtil.getFieldValue(root, childrenFieldName), handler, childrenFieldName);
		}
	}

	/**
	 * 树形结构迭代器（向下，深度优先）
	 *
	 * @param roots             根节点集合
	 * @param handler           处理器
	 * @param childrenExtractor 获取子节点集合的方法，如Map等结构只能通过get方法获取
	 * @param <T>               节点类型
	 */
	public static <T> void treeIteratorDown(Collection<T> roots, Consumer<T> handler, Function<T, Collection<T>> childrenExtractor) {
		if (CollUtil.isEmpty(roots)) {
			return;
		}
		for (T root : roots) {
			handler.accept(root);
			treeIteratorDown(childrenExtractor.apply(root), handler, childrenExtractor);
		}
	}

	/**
	 * 树形结构迭代器（向上）<br/>
	 * 从给定的node开始，向上遍历每个父节点，直到根节点<br/>
	 * 注意：只会一直向上寻找父节点，而不保证遍历所有节点<br/>
	 *
	 * @param node            开始遍历的节点
	 * @param handler         处理器
	 * @param parentFieldName 父节点字段名
	 * @param <T>             节点类型
	 */
	@SuppressWarnings("unchecked")
	public static <T> void treeIteratorUp(T node, Consumer<T> handler, String parentFieldName) {
		if (node == null) {
			return;
		}
		handler.accept(node);
		treeIteratorUp((T) ReflectUtil.getFieldValue(node, parentFieldName), handler, parentFieldName);
	}

	/**
	 * 从树中查找目标节点
	 *
	 * @param roots             根节点集合
	 * @param predicate         断言
	 * @param childrenFieldName 子节点字段名
	 * @param <T>               节点类型
	 * @return 目标节点
	 */
	@SuppressWarnings("unchecked")
	public static <T> Optional<T> findInTree(Collection<T> roots, Predicate<T> predicate, String childrenFieldName) {
		if (CollUtil.isEmpty(roots)) {
			return Optional.empty();
		}
		for (T root : roots) {
			if (predicate.test(root)) {
				return Optional.of(root);
			}
			Optional<T> optional = findInTree((Collection<T>) ReflectUtil.getFieldValue(root, childrenFieldName), predicate, childrenFieldName);
			if (optional.isPresent()) {
				return optional;
			}
		}
		return Optional.empty();
	}

	/**
	 * 从树中查找目标节点
	 *
	 * @param roots             根节点集合
	 * @param predicate         断言
	 * @param childrenExtractor 获取子节点集合的方法，如Map等结构只能通过get方法获取
	 * @param <T>               节点类型
	 * @return 目标节点
	 */
	public static <T> Optional<T> findInTree(Collection<T> roots, Predicate<T> predicate, Function<T, Collection<T>> childrenExtractor) {
		if (CollUtil.isEmpty(roots)) {
			return Optional.empty();
		}
		for (T root : roots) {
			if (predicate.test(root)) {
				return Optional.of(root);
			}
			Optional<T> optional = findInTree(childrenExtractor.apply(root), predicate, childrenExtractor);
			if (optional.isPresent()) {
				return optional;
			}
		}
		return Optional.empty();
	}

	/**
	 * 从树中查找目标节点，返回该节点的路径
	 *
	 * @param roots             根节点集合
	 * @param predicate         断言
	 * @param pathExtractor     获取路径元素的方法
	 * @param childrenExtractor 获取子节点集合的方法，如Map等结构只能通过get方法获取
	 * @param <T>               节点类型
	 * @param <P>               路径元素类型
	 * @return 路径
	 */
	public static <T, P> List<P> findPathToNode(Collection<T> roots, Predicate<T> predicate, Function<T, P> pathExtractor, Function<T, Collection<T>> childrenExtractor) {
		if (CollUtil.isEmpty(roots)) {
			return null;
		}
		for (T root : roots) {
			List<P> path = findPathToNode(root, predicate, pathExtractor, ArrayList::new, childrenExtractor);
			if (path != null) {
				return path;
			}
		}
		return null;
	}

	/**
	 * 从树中查找目标节点，返回该节点的路径
	 *
	 * @param root              根节点
	 * @param predicate         断言
	 * @param pathExtractor     获取路径元素的方法
	 * @param listFactory       路径列表工厂
	 * @param childrenExtractor 获取子节点集合的方法，如Map等结构只能通过get方法获取
	 * @param <T>               节点类型
	 * @param <P>               路径元素类型
	 * @return 路径
	 */
	public static <T, P> List<P> findPathToNode(T root, Predicate<T> predicate, Function<T, P> pathExtractor, Supplier<List<P>> listFactory, Function<T, Collection<T>> childrenExtractor) {
		if (listFactory == null) {
			throw new NullPointerException("路径列表工厂不能为空");
		}
		if (root == null) {
			throw new NullPointerException("节点不能为空");
		}
		List<P> path = listFactory.get();
		path.add(pathExtractor.apply(root));
		if (predicate.test(root)) {
			return path;
		}
		Collection<T> children = childrenExtractor.apply(root);
		if (CollUtil.isNotEmpty(children)) {
			for (T child : children) {
				List<P> findPath = findPathToNode(child, predicate, pathExtractor, listFactory, childrenExtractor);
				if (findPath != null) {
					path.addAll(findPath);
					return path;
				}
			}
		}
		return null;
	}


	/**
	 * 压平树
	 *
	 * @param roots             根节点集合
	 * @param childrenFieldName 子节点字段名
	 * @param <T>               节点类型
	 * @return 压平后的节点集合
	 */
	@SuppressWarnings("unchecked")
	public static <T> List<T> flattenTree(Collection<T> roots, String childrenFieldName) {
		List<T> list = new ArrayList<>();
		if (CollUtil.isEmpty(roots)) {
			return list;
		}
		for (T root : roots) {
			list.add(root);
			Collection<T> children = (Collection<T>) ReflectUtil.getFieldValue(root, childrenFieldName);
			if (!CollUtil.isEmpty(children)) {
				list.addAll(flattenTree(children, childrenFieldName));
			}
		}
		return list;
	}

	/**
	 * 压平树
	 *
	 * @param roots             根节点集合
	 * @param childrenFieldName 子节点字段名
	 * @param collector         收集器
	 * @param <T>               节点类型
	 * @param <A>               中间类型
	 * @param <R>               结果类型
	 * @return 压平后的节点集合
	 */
	public static <T, A, R> R flattenTree(Collection<T> roots, String childrenFieldName, Collector<? super T, A, R> collector) {
		List<T> flattenedTree = flattenTree(roots, childrenFieldName);
		return flattenedTree.stream().collect(collector);
	}

}
